from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark.ml.linalg import DenseVector
from pyspark.ml.feature import StandardScaler
from pyspark.ml.regression import LinearRegression

# 初始化 SparkSession 和 SparkContext
from pyspark.sql.functions import col
from pyspark.sql.types import FloatType

spark = SparkSession.builder.master("local").\
    appName("California Housing ").\
    config("spark.executor.memory", "1gb").\
    getOrCreate()

sc = spark.sparkContext

# 读取数据并创建 RDD
rdd = sc.textFile('cal_housing.data')

# 读取数据每个属性的定义并创建 RDD
header = sc.textFile('cal_housing.domain')


# collect() 函数会把所有数据都加载到内存中，如果数据很大的话，有可能会造成内存泄漏，所以要小心使用。
# 平时比较常见的方法是用 take() 函数去只读取 RDD 中的某几个元素。
header.collect()

##[u'longitude: continuous.', u'latitude: continuous.', u'housingMedianAge: continuous. ', u'totalRooms: continuous. ', u'totalBedrooms: continuous. ', u'population: continuous. ', u'households: continuous. ', u'medianIncome: continuous. ', u'medianHouseValue: continuous. ']


rdd.take(2)
# 各个属性之间用逗号分隔开来。这不利于我们之后的处理，因为我们可能会需要分别读取每个对象的各个属性。
# 所以，让我们用 map 函数把大字符串分隔成数组，这会方便我们的后续操作。

rdd = rdd.map(lambda line: line.split(","))
rdd.take(2)




def convertColumn(df, names, newType):
    for name in names:
        df = df.withColumn(name, df[name].cast(newType))
    return df

columns = ['households', 'housingMedianAge', 'latitude', 'longitude', 'medianHouseValue', 'medianIncome', 'population', 'totalBedRooms', 'totalRooms']


# DataFrame API 在查询结构化数据时更方便使用，而且性能更好。
# 在这个例子中你可以看到，数据的 schema 是定义好的，我们需要去查询各个列，
# 所以 DataFrame API 显然更加适用。所以，我们需要先把 RDD 转换为 DataFrame。

df = rdd.map(lambda line: Row(longitude=line[0],
                              latitude=line[1],
                              housingMedianAge=line[2],
                              totalRooms=line[3],
                              totalBedRooms=line[4],
                              population=line[5],
                              households=line[6],
                              medianIncome=line[7],
                              medianHouseValue=line[8])).toDF()

df = convertColumn(df, columns, FloatType())


df.groupBy("housingMedianAge").count().sort("housingMedianAge",ascending=False).show()


# 预处理
# 通过上面的数据分析，你可能会发现这些数据还是不够直观。具体的问题有：
#
# 1.房价的值普遍都很大，我们可以把它调整成相对较小的数字；
# 2.有的属性没什么意义，比如所有房子的总房间数和总卧室数，我们更加关心的是平均房间数；
# 3.在我们想要构建的线性模型中，房价是结果，其他属性是输入参数。所以我们需要把它们分离处理；
# 4.有的属性最小值和最大值范围很大，我们可以把它们标准化处理。


# 对于第一点，我们观察到大多数房价都是十万起的，所以可以用 withColumn() 函数把所有房价都除以 100000。
df = df.withColumn("medianHouseValue", col("medianHouseValue")/100000)


# 对于第二点，我们可以添加如下三个新的列：
# 每个家庭的平均房间数：roomsPerHousehold
# 每个家庭的平均人数：populationPerHousehold
# 卧室在总房间的占比：bedroomsPerRoom
df = df.withColumn("roomsPerHousehold", col("totalRooms")/col("households")).\
    withColumn("populationPerHousehold", col("population")/col("households")).\
    withColumn("bedroomsPerRoom", col("totalBedRooms")/col("totalRooms"))


df = df.select("medianHouseValue",
               "totalBedRooms",
               "population",
               "households",
               "medianIncome",
               "roomsPerHousehold",
               "populationPerHousehold",
               "bedroomsPerRoom")


# 对于第三点，最简单的办法就是把 DataFrame 转换成 RDD，然后用 map() 函数把每个对象分成两部分：房价和一个包含其余属性的列表，然后在转换回 DataFrame。
# 我们重新把两部分重新标记为“label”和“features”，label 代表的是房价，features 代表包括其余参数的列表。

input_data = df.rdd.map(lambda x: (x[0], DenseVector(x[1:])))
df = spark.createDataFrame(input_data, ["label", "features"])



# 对于第四点，数据的标准化我们可以借助 Spark 的机器学习库 Spark ML 来完成。Spark ML 也是基于 DataFrame，
# 它提供了大量机器学习的算法实现、
# 数据流水线（pipeline）相关工具和很多常用功能


# 1.特征选择
# 该行代码的结果是创建了一个 RDD input_data，其中每个元素都是包含房价和属性列表的元组。
input_data = df.rdd.map(lambda x: (x[0], DenseVector(x[1:])))
df = spark.createDataFrame(input_data, ["label", "features"])


# 2.特征变换
standardScaler = StandardScaler(inputCol="features", outputCol="features_scaled")
scaler = standardScaler.fit(df)
scaled_df = scaler.transform(df)

scaled_df.take(1)


# 这里我们将 80% 的数据用于训练，剩下 20% 作为测试集。
train_data, test_data = scaled_df.randomSplit([.8,.2],seed=123)


# 用 Spark ML 提供的 LinearRegression 功能，我们可以很容易得构建一个线性回归模型，如下所示。
lr = LinearRegression(featuresCol='features_scaled', labelCol="label", maxIter=10, regParam=0.3, elasticNetParam=0.8)
linearModel = lr.fit(train_data)



# 通过将测试数据 test_data 传递给线性模型 linearModel 的 transform 方法，
# 得到了对测试数据进行预测的结果。predicted 是一个包含预测结果的 DataFrame。
predicted = linearModel.transform(test_data)


# 3.特征提取
# 从 predicted DataFrame 中选择名称为 "prediction" 的列，并将其转换为 RDD。然后使用 map 方法通过 lambda 函数，
# 将 RDD 中的每个元素提取出来，构成一个包含预测值的 RDD。
predictions = predicted.select("prediction").rdd.map(lambda x: x[0])

# 从 predicted Dataframe 中选择名称为 "label" 的列，并将其转换为 RDD。然后使用 map 方法通过 lambda 函数，
# 将 RDD 中的每个元素提取出来，构成一个包含标签值的 RDD。
labels = predicted.select("label").rdd.map(lambda x: x[0])

# 将预测值的 RDD predictions 和标签值的 RDD labels 进行配对，使用 zip 方法将两个 RDD 中的元素按索引进行配对，
# 创建一个包含预测值和标签值配对的列表。然后使用 collect 方法收集这些配对，并将其返回给变量 predictionAndLabel。
predictionAndLabel = predictions.zip(labels).collect()

# 用 RDD 的 zip() 函数把预测值和真实值放在一起，这样可以方便地进行比较。比如让我们看一下前两个对比结果。
predictionAndLabel.take(2)
